"""
模型训练脚本
"""

import torch
import numpy as np
import logging
from predictor import QueueLengthPredictor
from config import SystemConfig
from utils import setup_logging

def generate_synthetic_data(num_samples=1000):
    """生成合成训练数据"""
    # 模拟排队长度数据
    times = np.arange(num_samples)
    
    # 基础模式：早晚高峰
    morning_peak = 0.7 * np.exp(-0.5 * ((times % 1000 - 300) / 100) ** 2)
    evening_peak = 0.8 * np.exp(-0.5 * ((times % 1000 - 700) / 120) ** 2)
    
    # 随机波动
    noise = 0.1 * np.random.randn(num_samples)
    
    # 合成排队长度
    queue_lengths = 20 + 60 * (morning_peak + evening_peak) + noise
    queue_lengths = np.clip(queue_lengths, 0, 100)
    
    return queue_lengths

def main():
    """主训练函数"""
    setup_logging()
    logger = logging.getLogger(__name__)
    
    # 加载配置
    config = SystemConfig()
    
    # 初始化预测器
    predictor = QueueLengthPredictor(config.prediction)
    
    # 生成训练数据
    logger.info("生成训练数据...")
    synthetic_data = generate_synthetic_data(5000)
    
    # 训练模型
    logger.info("开始训练模型...")
    best_loss = float('inf')
    patience = 10
    patience_counter = 0
    
    for epoch in range(100):
        # 更新历史数据
        for length in synthetic_data[epoch * 50: (epoch + 1) * 50]:
            predictor.update_history(length)
        
        # 训练一个周期
        loss = predictor.train_epoch()
        
        logger.info(f"Epoch {epoch + 1}, Loss: {loss:.4f}")
        
        # 早停检查
        if loss < best_loss:
            best_loss = loss
            patience_counter = 0
            predictor.save_model("models/best_predictor.pth")
        else:
            patience_counter += 1
            
        if patience_counter >= patience:
            logger.info("早停触发")
            break
    
    # 保存最终模型
    predictor.is_trained = True
    predictor.save_model()
    logger.info("训练完成，模型已保存")

if __name__ == "__main__":
    main()